#include "nvds_sourcebin.h"
#include "nvds_custom_meta.h"
#include <opencv2/opencv.hpp>
#include <spdlog/spdlog.h>
#include <stdexcept>
#include <nvbufsurftransform.h>

#define GST_CAPS_FEATURES_NVMM "memory:NVMM"

static uint64_t getCurrentMilliseconds()
{
    auto tp = std::chrono::high_resolution_clock::now().time_since_epoch();
    return std::chrono::duration_cast<std::chrono::milliseconds>(tp).count();
}

NVDSSourcebin::NVDSSourcebin(
    const char *bin_name, const char *uri, bool use_transform, int img_quality, bool scale, int width, int height) :
    m_surfaceTransform(use_transform),
    m_jpgencQuality(img_quality),
    m_jpgencScaleEnable(scale),
    m_jpgencScaleWidth(width),
    m_jpgencScaleHeight(height)
{
    m_sourcebin = createSourcebin(bin_name, uri);
    if (!m_sourcebin)
        throw std::runtime_error("Create sourcebin failed!");
}

NVDSSourcebin::~NVDSSourcebin()
{
}

void NVDSSourcebin::padAddedCallback(GstElement *decodebin, GstPad *decoder_src_pad, gpointer data)
{
    GstCaps *caps = gst_pad_get_current_caps(decoder_src_pad);
    if (!caps) {
        caps = gst_pad_query_caps(decoder_src_pad, NULL);
    }
    const GstStructure *str        = gst_caps_get_structure(caps, 0);
    const gchar        *name       = gst_structure_get_name(str);
    GstElement         *source_bin = (GstElement *)data;
    GstCapsFeatures    *features   = gst_caps_get_features(caps, 0);

    if (strncmp(name, "video", 5) == 0) {
        if (gst_caps_features_contains(features, GST_CAPS_FEATURES_NVMM)) {
            GstPad *bin_ghost_pad = gst_element_get_static_pad(source_bin, "src");
            if (!gst_ghost_pad_set_target(GST_GHOST_PAD(bin_ghost_pad), decoder_src_pad)) {
                g_printerr("Failed to link decoder src pad to source bin ghost pad\n");
            }
            gst_object_unref(bin_ghost_pad);
        } else {
            g_printerr("Error: Decodebin did not pick nvidia decoder plugin.\n");
        }
    }
}

void NVDSSourcebin::childAddedCallback(GstChildProxy *child_proxy, GObject *object, gchar *name, gpointer user_data)
{
    if (g_strrstr(name, "decodebin") == name) {
        g_signal_connect(G_OBJECT(object), "child-added", G_CALLBACK(childAddedCallback), user_data);
    }
    if (g_strrstr(name, "src") == name && g_strrstr(name, "src_") != name) {
        g_object_set(G_OBJECT(object), "drop-on-latency", true, NULL);
        g_object_set(G_OBJECT(object), "latency", 500, NULL);
        // 指定 TCP 方式进行连接
        g_object_set(G_OBJECT(object), "protocols", 0x04, NULL);
    }

    if (g_strrstr(name, "nvv4l2decoder") == name) {
#ifdef __aarch64__
        g_object_set(G_OBJECT(object), "disable-dpb", true, NULL);
#endif
    }
}

void NVDSSourcebin::surfaceTransform(
    NVDSCustomMeta *user_metadata, NvBufSurface *surface, int index, int quality, bool scale, int scale_width, int scale_height)
{
    const NvBufSurfaceParams &params = surface->surfaceList[index];

    // 设置转换参数
    NvBufSurfTransformParams transform_params;
    NvBufSurfTransformRect   src_rect = {.top = 0, .left = 0, .width = params.width, .height = params.height};
    NvBufSurfTransformRect   dst_rect = {.top = 0, .left = 0, .width = params.width, .height = params.height};
    transform_params.src_rect         = &src_rect;
    transform_params.dst_rect         = &dst_rect;
    transform_params.transform_flag   = NVBUFSURF_TRANSFORM_CROP_SRC | NVBUFSURF_TRANSFORM_CROP_DST;
    transform_params.transform_filter = NvBufSurfTransformInter_Default;

    // 创建 RGB surface, 数据转换
    NvBufSurfaceCreateParams surface_create_params;
    surface_create_params.gpuId       = surface->gpuId;
    surface_create_params.width       = params.width;
    surface_create_params.height      = params.height;
    surface_create_params.size        = 0;
    surface_create_params.colorFormat = NVBUF_COLOR_FORMAT_BGRA;
    surface_create_params.layout      = NVBUF_LAYOUT_PITCH;
    surface_create_params.memType     = NVBUF_MEM_DEFAULT;

    NvBufSurface *rgb_surface = NULL;
    NvBufSurfaceCreate(&rgb_surface, surface->batchSize, &surface_create_params);
    if (NvBufSurfTransform(surface, rgb_surface, &transform_params) != NvBufSurfTransformError_Success) {
        spdlog::error("NvBufSurfTransform failed with error while converting buffer.");
        return;
    }

    const NvBufSurfaceParams &rgb_params = rgb_surface->surfaceList[index];
    // 映射 GPU 内存到 CPU
    if (NvBufSurfaceMap(rgb_surface, index, -1, NVBUF_MAP_READ) != 0) {
        spdlog::error("NvBuf surface map failed!");
        NvBufSurfaceDestroy(rgb_surface);
        return;
    }
    NvBufSurfaceSyncForCpu(rgb_surface, 0, 0);

    // 将帧数据转换为 OpenCV 的 Mat 格式
    auto    stime = std::chrono::system_clock::now();
    cv::Mat image_mat(rgb_params.height, rgb_params.width, CV_8UC4, rgb_params.mappedAddr.addr[0]);
    if (scale) {
        cv::resize(image_mat, image_mat, cv::Size(scale_width, scale_height));
    }

    // 图片转换为 jpg
    std::vector<int> opencv_params({cv::IMWRITE_JPEG_QUALITY, quality});
    if (!cv::imencode(".jpg", image_mat, user_metadata->imgBuffer(), opencv_params)) {
        spdlog::error("surfaceTransform failed!");
        NvBufSurfaceUnMap(rgb_surface, index, -1); // 取消映射
        NvBufSurfaceDestroy(rgb_surface);
        return;
    }

    if (scale)
        user_metadata->setImageSize(scale_width, scale_height);
    else
        user_metadata->setImageSize(rgb_params.width, rgb_params.height);

    // 计算耗时
    auto etime  = std::chrono::system_clock::now();
    auto costms = std::chrono::duration<double, std::ratio<1, 1000>>(etime - stime).count();
    spdlog::debug("jpeg encode cost {}ms.", costms);

    NvBufSurfaceUnMap(rgb_surface, index, -1); // 取消映射
    NvBufSurfaceDestroy(rgb_surface);
}

GstPadProbeReturn NVDSSourcebin::probeCallback(GstPad *pad, GstPadProbeInfo *info, gpointer u_data)
{
    uint64_t        curtime                  = getCurrentMilliseconds();
    NVDSSourcebin  *pthis                    = (NVDSSourcebin *)u_data;
    GstBuffer      *buf                      = GST_PAD_PROBE_INFO_BUFFER(info);
    NvBufSurface   *surface                  = NULL;
    GstMapInfo      map_info                 = GST_MAP_INFO_INIT;
    NVDSCustomMeta *user_metadata            = new NVDSCustomMeta();
    user_metadata->timestamp()[LANDING_TIME] = curtime;

    if (pthis->m_surfaceTransform) {
        if (gst_buffer_map(buf, &map_info, GST_MAP_READ)) {
            surface = (NvBufSurface *)map_info.data;
            if (surface) {
                surfaceTransform(user_metadata,
                                 surface,
                                 0,
                                 pthis->m_jpgencQuality,
                                 pthis->m_jpgencScaleEnable,
                                 pthis->m_jpgencScaleWidth,
                                 pthis->m_jpgencScaleHeight);
            }
        }
        gst_buffer_unmap(buf, &map_info);
    }

    user_metadata->gstBufferAddMeta(buf);
    return GST_PAD_PROBE_OK;
}

GstElement *NVDSSourcebin::createSourcebin(const char *bin_name, const char *uri)
{
    GstElement *uri_decode_bin = NULL;
    GstElement *bin            = gst_bin_new(bin_name);

    /** uridecodebin 是一个通用的解码元素，适用于各种不同的媒体流格式，而 nvurisrcbin 是特定于 NVIDIA 的，可能只适用于 NVIDIA
     * 设备或特定协议的数据流。*/
    // uri_decode_bin = gst_element_factory_make("uridecodebin", "uri-decode-bin");
    uri_decode_bin = gst_element_factory_make("nvurisrcbin", "uri-decode-bin");
    // g_object_set(G_OBJECT(uri_decode_bin), "file-loop", false, NULL);
    // g_object_set(G_OBJECT(uri_decode_bin), "cudadec-memtype", 0, NULL);

    if (!bin || !uri_decode_bin) {
        g_printerr("One element in source bin could not be created.\n");
        return NULL;
    }

    g_object_set(G_OBJECT(uri_decode_bin), "uri", uri, NULL);
    if (strncmp(uri, "rtsp:", 5) == 0) {
        // g_object_set(G_OBJECT(uri_decode_bin), "rtsp-reconnect-attempts", -1, NULL);
        g_object_set(G_OBJECT(uri_decode_bin), "rtsp-reconnect-interval", 10, NULL); // 设置重连
    }

    // 控制输出帧数，设置多少帧图片里保存1帧
    g_object_set(G_OBJECT(uri_decode_bin), "drop-frame-interval", 1, NULL);
    g_signal_connect(G_OBJECT(uri_decode_bin), "pad-added", G_CALLBACK(padAddedCallback), bin);
    g_signal_connect(G_OBJECT(uri_decode_bin), "child-added", G_CALLBACK(childAddedCallback), bin);

    gst_bin_add(GST_BIN(bin), uri_decode_bin);
    if (!gst_element_add_pad(bin, gst_ghost_pad_new_no_target("src", GST_PAD_SRC))) {
        g_printerr("Failed to add ghost pad in source bin\n");
        return NULL;
    }
    return bin;
}

void NVDSSourcebin::bindStreammux(const char *pad_name, GstElement *streammux)
{
    GstPad *srcpad  = gst_element_get_static_pad(m_sourcebin, "src");
    GstPad *sinkpad = gst_element_get_request_pad(streammux, pad_name);
    if (!srcpad || !sinkpad) {
        throw std::runtime_error("Get sourcebin pad failed!");
    }
    if (gst_pad_link(srcpad, sinkpad) != GST_PAD_LINK_OK) {
        throw std::runtime_error("Failed to link decoder to stream muxer.");
    }

    gst_pad_add_probe(srcpad, GST_PAD_PROBE_TYPE_BUFFER, probeCallback, this, NULL);

    gst_object_unref(srcpad);
    gst_object_unref(sinkpad);
}

void NVDSSourcebin::bindPipeline(GstElement *pipeline)
{
    gst_bin_add(GST_BIN(pipeline), m_sourcebin);
}
